{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "from pyskl.smp import *\n",
    "from pyskl.utils.visualize import Vis3DPose, Vis2DPose\n",
    "from mmcv import load, dump"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download annotations\n",
    "download_file('http://download.openmmlab.com/mmaction/pyskl/demo/annotations/ntu60_samples_hrnet.pkl', 'ntu60_2d.pkl')\n",
    "download_file('http://download.openmmlab.com/mmaction/pyskl/demo/annotations/ntu60_samples_3danno.pkl', 'ntu60_3d.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize 2D Skeletons without video\n",
    "annotations = load('ntu60_2d.pkl')\n",
    "index = 0\n",
    "anno = annotations[index]\n",
    "vid = Vis2DPose(anno, thre=0.2, out_shape=(540, 960), layout='coco', fps=12, video=None)\n",
    "vid.ipython_display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize 2D Skeletons with the original RGB Video\n",
    "annotations = load('ntu60_2d.pkl')\n",
    "index = 0\n",
    "anno = annotations[index]\n",
    "frame_dir = anno['frame_dir']\n",
    "video_url = f\"http://download.openmmlab.com/mmaction/pyskl/demo/nturgbd/{frame_dir}.avi\"\n",
    "download_file(video_url, frame_dir + '.avi')\n",
    "vid = Vis2DPose(anno, thre=0.2, out_shape=(540, 960), layout='coco', fps=12, video=frame_dir + '.avi')\n",
    "vid.ipython_display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize 3D Skeletons\n",
    "from pyskl.datasets.pipelines import PreNormalize3D\n",
    "annotations = load('ntu60_3d.pkl')\n",
    "index = 0\n",
    "anno = annotations[index]\n",
    "anno = PreNormalize3D()(anno)  # * Need Pre-Normalization before Visualization\n",
    "vid = Vis3DPose(anno, layout='nturgb+d', fps=12, angle=(30, 45), fig_size=(8, 8), with_grid=False)\n",
    "vid.ipython_display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean directories\n",
    "os.remove('ntu60_3d.pkl')\n",
    "os.remove('ntu60_2d.pkl')\n",
    "for f in glob.glob(\"S*.avi\"):\n",
    "    os.remove(f)\n",
    "os.remove('__temp__.mp4')"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "8516388d98336f8f25b57d66898c29467d9fc2c465b9ee1e16d04a4702b7bff1"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
